import torch
import torch.nn as nn
import torch.optim as optim
from dataset import get_data_loaders
import time
import os
from models import ConvNet, ModifiedAlexNet, ModifiedResNet, FCN

def train_model(model, train_loader, test_loader, num_epochs=10, learning_rate=0.001):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    # 创建保存模型的目录
    if not os.path.exists('checkpoints'):
        os.makedirs('checkpoints')
    
    best_accuracy = 0.0
    
    for epoch in range(num_epochs):
        start_time = time.time()
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        # 训练阶段
        for i, (images, labels) in enumerate(train_loader):
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            
            if (i + 1) % 100 == 0:
                print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], '
                      f'Loss: {running_loss/100:.4f}, '
                      f'Accuracy: {100 * correct/total:.2f}%')
                running_loss = 0.0
        
        # 测试阶段
        model.eval()
        test_correct = 0
        test_total = 0
        
        with torch.no_grad():
            for images, labels in test_loader:
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                test_total += labels.size(0)
                test_correct += (predicted == labels).sum().item()
        
        test_accuracy = 100 * test_correct / test_total
        epoch_time = time.time() - start_time
        
        print(f'Epoch [{epoch+1}/{num_epochs}] completed in {epoch_time:.2f}s')
        print(f'Test Accuracy: {test_accuracy:.2f}%')
        
        # 保存最佳模型
        if test_accuracy > best_accuracy:
            best_accuracy = test_accuracy
            torch.save(model.state_dict(), 'checkpoints/best_model.pth')
            print(f'保存最佳模型，准确率: {best_accuracy:.2f}%')
        
        # 每个epoch都保存一次
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': running_loss,
            'accuracy': test_accuracy
        }, f'checkpoints/epoch_{epoch+1}.pth')

def main():
    # 数据加载器
    train_loader, test_loader, num_classes = get_data_loaders(
        batch_size=64,
        num_workers=0
    )
    
    # 创建模型
    model_name = "alexnet"  # 建议使用更复杂的模型来处理多语言字符
    
    # 创建模型
    if model_name == "convnet":
        model = ConvNet(num_classes=num_classes)
    elif model_name == "alexnet":
        model = ModifiedAlexNet(num_classes=num_classes)
    elif model_name == "resnet":
        model = ModifiedResNet(num_classes=num_classes)
    elif model_name == "fcn":
        model = FCN(num_classes=num_classes)
    else:
        raise ValueError(f"未知的模型名称: {model_name}")
    
    # 训练模型
    train_model(
        model=model,
        train_loader=train_loader,
        test_loader=test_loader,
        num_epochs=10,
        learning_rate=0.001
    )

if __name__ == '__main__':
    main() 